import dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

dotenv.load_dotenv()

def retrieval(query: str) -> str:
    """一个模拟的检索器"""
    print("正在检索", query)
    return "我是慕小课"

# 1.编排prompt
prompt = ChatPromptTemplate.from_template("""请根据用户的问题回答，可以参考对应的上下文进行生成。
<context>
{context}
</context>
用户的提问是：{query}
""")

# 2.构建大语言模型
llm = ChatOpenAI(model_name="kimi-k2-0711-preview")

# 3.创建输出解析器
parser = StrOutputParser()

# 4.编排链
chain = RunnablePassthrough.assign(context=lambda x: retrieval(x["query"])) | prompt | llm | parser

# 5.调用链得到结果
res = chain.invoke({"query": "你好，大姐"})

print(res, end='')
